# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
import random
import torch
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


SHAPE_LIST = [[[9728, 7168], [4, 128, 448, 16, 32], [4, 4096], [9728], [4]],
 [[8192, 7168], [4, 64, 448, 16, 32], [4, 2048], [8192], [4]],
 [[512, 7168], [4, 128, 448, 16, 32], [4, 4096], [512], [4]],
 [[536, 7168], [160, 128, 448, 16, 32], [160, 4096], [536], [160]],
 [[16384, 7168], [4, 128, 448, 16, 32], [4, 4096], [16384], [4]],
 [[25856, 7168], [4, 128, 448, 16, 32], [4, 4096], [25856], [4]]
 ]

@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul_swiglu_quant_model")
class AscendGroupedMatmulSwigluQuant(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)


    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        x = case_config.inputs[0]
        weight = case_config.inputs[1]
        weightScale = case_config.inputs[4]
        xScale = case_config.inputs[5]
        groupList = case_config.inputs[6]

        case: int = case_config.inputs[8].range_values
        try:
            x.shape, weight.shape, weightScale.shape, xScale.shape, groupList.shape = SHAPE_LIST[case]
        except:
            print(case)
            exit()
        return case_config